#! /usr/bin/env python
# coding=utf-8

import argparse
import sys
from argparse import RawTextHelpFormatter
import os
from Bio import SeqIO
from ete3 import PhyloTree
import re


parser = argparse.ArgumentParser(
    description='''
    读取cafe的结果 这里物种名称千万不能写下划线 _ 物种名称使用逗号分割 如 Sgua, Gleg
    read_cafe_and_give_result.py -c lambdaresult1.cafe -r Orthogroups.GeneCount_filtered.tsv -a Sgua, Gleg -b Gfla -o out.tsv

    ''',formatter_class=RawTextHelpFormatter)



parser.add_argument('-c',
                help='cafe 结果')

parser.add_argument('-r',
                help='用于cafe分析的OGS结果')

parser.add_argument('-a',
                help='目标物种列表 用,分割')

parser.add_argument('-b',
                help='对照物种列表 用,分割')

parser.add_argument('-o',
                help='输出结果')



args = parser.parse_args()


if not args.c or not args.r or not args.a or not args.b or not args.o:
    parser.print_help()
    sys.exit()


infile = args.c

ogs_list_file = args.r 


outfile = args.o

select_species = [i.strip() for i in args.a.split(',')]

compare_species = [i.strip() for i in args.b.split(',')]


handel = open(infile)


expanded_lista = [] 
contract_lista = []



def get_target_node(tree,lista=[]):
	# 判断这个节点是不是一个 包含我们需要物种的节点
	k = 1
	target_node = []
	node_name = tree.name.split('_')[-1]
	for i in tree.get_leaf_names():
		spe_name = i.split('_')[0]
		if spe_name not in lista:
			k = 0
			break
	if k ==1:
		target_node.append(node_name)

	for i in tree.get_children():
		new_node  = get_target_node(i,lista)

		#if new_node != []:
		target_node = new_node + target_node

	return target_node


def get_target_leafnode(tree,lista=[]):
	# 判断这个节点是不是一个 包含我们需要物种的节点
	k = 1
	target_node = []
	node_name = tree.name.split('_')[-1]
	for i in tree.get_leaf_names():
		spe_name = i.split('_')[0]
		node_name = i.split('_')[1]
		if spe_name in lista:
			target_node.append(node_name)
	return target_node




# 读取ogslist 文件
ogs_dic = {}
with open(ogs_list_file) as fila:
	head_line = next(fila).strip().split('\t')
	for i in fila:
		k = i.strip().split('\t')
		if len(k)>2:
			ogs_dic[k[1]] = (k[0],k)

# 对输出文件 写入 head
outfile = open(outfile,'w')
outfile.write('\t'.join(head_line)+'\t'+'expand or contract'+'\t'+'Family-wide P-value'+'\t'+'Newick'+'\t'+	'Viterbi P-values'+'\n')

for i in handel:
	i = i.strip()
	if i.find('# IDs of nodes:')!=-1:
		tree_with_id_node = i.split(':')[1].replace(')<',')').replace('<','_').replace('>','')+';'
		break

t = PhyloTree(tree_with_id_node,format =1)


# 获取需要读取的 进行比较的节点
target_node = get_target_node(t,select_species)
compare_node = get_target_node(t,compare_species)

# 获得目标物种的节点
target_spe_node = get_target_leafnode(t,select_species)
compare_spe_node = get_target_leafnode(t,compare_species)


replace_re=  re.compile(r':[\d|.]+')

for i in handel:
	if i.find('(node ID, node ID): ')!=-1:
		node_order = [j.strip() for j in i.split('(node ID, node ID): ')[1].strip().replace('(','').replace(')','').replace(',',' ').split(' ')]
		#print(node_order)
		break

for i in handel:
	if i.find('\'ID\'')!=-1:
		break



for i in handel:
	k = i.strip().split('\t')

	
	if len(k)>=4 and float(k[2]) < 0.05:
	

		ogs_order = k[0]	
		gene_nub_tree = PhyloTree(replace_re.sub('',k[1].replace(')_',')'))+';',format =8)
		target_value = [int(j) for j in  get_target_node(gene_nub_tree,select_species)]
		compare_value = [int(j) for j in get_target_node(gene_nub_tree,compare_species)]
		target_spe_value = [int(j) for j in  get_target_leafnode(gene_nub_tree,select_species)]
		compare_spe_value = [int(j) for j in  get_target_leafnode(gene_nub_tree,compare_species)]

		# 需要在 选择分支中存在 显著的 值
		node_p_value_lista = k[3].strip().replace('(','').replace(')','').replace(',',' ').split(' ')
		node_p_value_dic = {x:y for x,y in zip(node_order,node_p_value_lista)}

		significant = 0
		for target in target_node:
	
			if float(node_p_value_dic[target])<0.05:
				significant = 1
				break

		line_lista = ogs_dic[k[0]][1]
		if significant==1:


			if min(target_spe_value)>max(compare_spe_value):
				expanded_lista.append(ogs_order)
				outfile.write('\t'.join(line_lista)+'\texpansion\t'+k[2]+'\t'+k[1]+'\t'+k[3]+'\n')

			if max(target_spe_value) < min(compare_spe_value):
				contract_lista.append(ogs_order)
				outfile.write('\t'.join(line_lista)+'\tcontraction\t'+k[2]+'\t'+k[1]+'\t'+k[3]+'\n')





outfile.close()





node_id_dic = {}


